#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
期货品种交易所代码查询工具
输入品种名称，输出该品种的当前合约交易所代码
"""

import sys
import time
from datetime import datetime, date
from tqsdk import TqApi, TqAuth
from rich.console import Console
from rich.table import Table
from rich.prompt import Prompt

# 初始化控制台输出
console = Console()

def get_futures_symbols(api, symbol_name):
    """
    根据品种名称获取相关的期货合约信息
    
    Args:
        api: TqApi实例
        symbol_name: 品种名称（如：cj, y, rb等）
    
    Returns:
        list: 包含合约信息的列表
    """
    try:
        # 获取所有可交易的合约
        all_symbols = api.query_quotes(ins_class="FUTURE")
        
        # 过滤出包含指定品种名称的合约
        filtered_symbols = [s for s in all_symbols if symbol_name.lower() in s.lower()]
        
        # 按合约代码排序
        filtered_symbols.sort()
        
        return filtered_symbols
    except Exception as e:
        console.print(f"获取合约信息失败: {e}", style="red")
        return []

def get_symbol_details(api, symbol):
    """
    获取合约的详细信息
    
    Args:
        api: TqApi实例
        symbol: 合约代码
    
    Returns:
        dict: 合约详细信息
    """
    try:
        # 获取合约的quote信息
        quote = api.get_quote(symbol)
        
        details = {
            'symbol': symbol,
            'instrument_name': quote.get('instrument_name', 'N/A'),
            'exchange_id': quote.get('exchange_id', 'N/A'),
            'product_id': quote.get('product_id', 'N/A'),
            'expire_date': quote.get('expire_date', 'N/A'),
            'trading_day': quote.get('trading_day', 'N/A'),
            'pre_close': quote.get('pre_close', 'N/A'),
            'last_price': quote.get('last_price', 'N/A'),
            'volume': quote.get('volume', 'N/A'),
            'open_interest': quote.get('open_interest', 'N/A')
        }
        
        return details
    except Exception as e:
        console.print(f"获取合约 {symbol} 详情失败: {e}", style="yellow")
        return None

def is_contract_expired(expire_date):
    """
    检查合约是否已到期
    
    Args:
        expire_date: 合约到期日
    
    Returns:
        bool: True表示已到期，False表示未到期
    """
    if expire_date == 'N/A':
        return False  # 无法获取到期日，默认未到期
    
    try:
        # 将到期日转换为日期对象
        if isinstance(expire_date, str):
            expire_date = datetime.strptime(expire_date[:10], "%Y-%m-%d").date()
        elif isinstance(expire_date, datetime):
            expire_date = expire_date.date()
        
        # 获取当前日期
        current_date = date.today()
        
        # 如果到期日早于当前日期，则已到期
        return expire_date < current_date
    except Exception as e:
        console.print(f"检查合约到期日失败: {e}", style="yellow")
        return False

def filter_active_contracts(symbols_details):
    """
    过滤出未到期的活跃合约
    
    Args:
        symbols_details: 所有合约详细信息列表
    
    Returns:
        list: 未到期的合约列表
    """
    active_contracts = []
    expired_contracts = []
    
    for details in symbols_details:
        if details:
            if is_contract_expired(details['expire_date']):
                expired_contracts.append(details)
            else:
                active_contracts.append(details)
    
    if expired_contracts:
        console.print(f"📅 已过滤掉 {len(expired_contracts)} 个已到期合约", style="yellow")
    
    return active_contracts

def display_symbols_table(symbols_details, show_only_active=True):
    """
    以表格形式显示合约信息
    
    Args:
        symbols_details: 合约详细信息列表
        show_only_active: 是否只显示未到期合约
    """
    if not symbols_details:
        console.print("未找到相关合约信息", style="yellow")
        return
    
    # 如果需要只显示未到期合约，先进行过滤
    if show_only_active:
        symbols_details = filter_active_contracts(symbols_details)
        if not symbols_details:
            console.print("❌ 未找到未到期的活跃合约", style="red")
            return
    
    table_title = "期货合约信息查询结果" + (" (仅显示未到期合约)" if show_only_active else "")
    table = Table(title=table_title, show_header=True, header_style="bold magenta")
    
    # 添加表头
    table.add_column("合约代码", style="cyan", width=15)
    table.add_column("合约名称", style="green", width=20)
    table.add_column("交易所", style="blue", width=8)
    table.add_column("品种代码", style="yellow", width=10)
    table.add_column("到期日", style="white", width=12)
    table.add_column("最新价", style="red", width=10)
    table.add_column("成交量", style="white", width=10)
    table.add_column("持仓量", style="white", width=10)
    
    # 添加数据行
    for details in symbols_details:
        if details:
            # 检查合约是否已到期，为已到期合约添加特殊标记
            is_expired = is_contract_expired(details['expire_date'])
            symbol_style = "dim" if is_expired else "cyan"
            
            table.add_row(
                f"{details['symbol']}{' (已到期)' if is_expired else ''}",
                details['instrument_name'],
                details['exchange_id'],
                details['product_id'],
                str(details['expire_date'])[:10] if details['expire_date'] != 'N/A' else 'N/A',
                str(details['last_price']) if details['last_price'] != 'N/A' else 'N/A',
                str(details['volume']) if details['volume'] != 'N/A' else 'N/A',
                str(details['open_interest']) if details['open_interest'] != 'N/A' else 'N/A'
            )
    
    console.print(table)

def main():
    """主函数"""
    console.print("🎯 期货品种交易所代码查询工具", style="bold green")
    console.print("=" * 60, style="bold blue")
    
    # 初始化TqApi
    try:
        api = TqApi(auth=TqAuth("cps168", "alibaba"))
        console.print("✅ 天勤量化API连接成功", style="green")
    except Exception as e:
        console.print(f"❌ 天勤量化API连接失败: {e}", style="red")
        return
    
    # 询问用户是否只显示未到期合约
    show_only_active = Prompt.ask(
        "是否只显示未到期合约？",
        choices=["Y", "N"],
        default="Y"
    )
    show_only_active = show_only_active == "Y"
    
    if show_only_active:
        console.print("📅 将只显示未到期的活跃合约", style="bold yellow")
    else:
        console.print("📊 将显示所有合约（包括已到期合约）", style="bold yellow")
    
    while True:
        console.print("\n" + "-" * 40, style="bold blue")
        
        # 获取用户输入的品种名称
        symbol_name = Prompt.ask(
            "请输入期货品种名称",
            default="cj",
            choices=["cj", "y", "rb", "i", "CF", "ao", "p", "jm", "j", "退出"]
        )
        
        if symbol_name == "退出":
            break
        
        console.print(f"🔍 正在查询品种 '{symbol_name}' 的相关合约...", style="bold yellow")
        
        # 获取合约列表
        symbols = get_futures_symbols(api, symbol_name)
        
        if not symbols:
            console.print(f"❌ 未找到品种 '{symbol_name}' 的相关合约", style="red")
            continue
        
        console.print(f"✅ 找到 {len(symbols)} 个相关合约", style="green")
        
        # 获取每个合约的详细信息
        symbols_details = []
        for symbol in symbols:
            details = get_symbol_details(api, symbol)
            if details:
                symbols_details.append(details)
            # 避免请求过快
            time.sleep(0.1)
        
        # 显示结果
        display_symbols_table(symbols_details, show_only_active=show_only_active)
        
        # 显示当前主力合约（通常是最活跃的合约）
        if symbols_details:
            # 过滤出未到期合约用于主力合约判断
            active_contracts = filter_active_contracts(symbols_details) if show_only_active else symbols_details
            
            if active_contracts:
                # 按成交量排序，找出主力合约
                active_symbols = [s for s in active_contracts if s['volume'] != 'N/A']
                if active_symbols:
                    active_symbols.sort(key=lambda x: x['volume'] if x['volume'] != 'N/A' else 0, reverse=True)
                    main_contract = active_symbols[0]
                    
                    # 检查主力合约是否已到期
                    is_expired = is_contract_expired(main_contract['expire_date'])
                    status_text = " (已到期)" if is_expired else ""
                    
                    console.print(f"\n🏆 当前主力合约: [bold cyan]{main_contract['symbol']}[/bold cyan]{status_text} "
                                f"(成交量: {main_contract['volume']})", style="bold green")
    
    # 关闭API连接
    try:
        api.close()
        console.print("\n👋 程序已退出，感谢使用！", style="bold green")
    except:
        pass

if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        console.print("\n👋 程序被用户中断", style="yellow")
    except Exception as e:
        console.print(f"\n❌ 程序运行出错: {e}", style="red")